from gym import spaces, Env
import numpy as np
import time
from bdtime import tt

import gym
import gym.spaces
from typing import Generic, Optional, SupportsFloat, Tuple, TypeVar, Union
from gym.utils import seeding


ObsType = TypeVar("ObsType")
ActType = TypeVar("ActType")


class TestEnv(Env):
    metadata = {
        'render.modes': ['human', 'rgb_array', 'custom'],
        'video.frames_per_second': 2
    }

    def __init__(self, l=10):
        self.l = l      # 左右路程线段的总长度
        self.action_space = spaces.Discrete(3)  # 0, 1, 2: 不动，左右
        # self.observation_space = spaces.Discrete(l)
        self.observation_space: gym.spaces.Box = gym.spaces.Box(np.array([-1]), np.array([l]))
        # self.observation_space = spaces.Box(-l, l, dtype=np.float32)
        self.state = None

        self._msg = ""
        self.time_delta = 0.3

    def step(self, action: int):
        assert self.action_space.contains(action), "{} {} invalid".format(
            action, type(action))
        if action == 1:
            self.state = max(0, self.state - 1)
            reward = 0
        elif action == 2:
            if self.state == self.l - 2:
                reward = 1
            else:
                reward = 0
            self.state = min(self.l - 1, self.state + 1)
        else:
            reward = 0
        if self.state == self.l - 1 or self.state == 0:
            done = True
        else:
            done = False

        ret = np.array([self.state], dtype=np.float32), reward, done, {}
        return ret
        # return self.state, reward, done, {}

    def reset(self, *args, **kwargs):
        self.state = np.random.choice(range(1,self.l-1))
        # return self.state
        return np.array([self.state], dtype=np.float32)

    def render(self, mode='human'):
        if mode == 'human':
            for i in range(self.l):
                if i == self.state:
                    print('X', end=' ')
                else:
                    print('_', end=' ')
            print()
            time.sleep(self.time_delta)
        elif mode == 'custom':
            msg = ''
            _len = len(self._msg)
            for i in range(self.l):
                if i == self.state:
                    msg += "X"
                    # print('X', end=' ')
                else:
                    msg += "_"
                    # print('_', end=' ')
                msg += ' '

            # ret = '\b' * _len + msg
            # print('\b' * _len + msg, end="")
            # _len = len(msg)
            # self._msg = msg
            time.sleep(self.time_delta)
            return msg

            # for i in range(5):
            #     tt.sleep(0.3)
            #     msg = f"---{tt.now(3)}"
            #     print('\b' * _len + msg, end="")
            #     _len = len(msg)

    def close(self):
        return None


if __name__ == '__main__':
    env = TestEnv(10)
    s = env.reset()
    while True:
        env.render()
        # u = RL.choose_action(s)
        s_, r_, done, _ = env.step(2)
        # RL.learn(s, u, r_, s_)
        if done:
            break
        s = s_
